package com.lu.utils;

import com.alibaba.fastjson.JSONObject;
import io.searchbox.client.JestClient;
import io.searchbox.client.JestResult;
import io.searchbox.core.*;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 * @program Elasticsearch
 * @description: es jest 查询 工具
 * @author: zhanglu
 * @create: 2020-04-30 08:50:00
 */
@Setter
@Component
@Slf4j
public class JestSearchUtil {

    @Autowired
    private JestClient jestClient;

    /******************************************普通查询******************************************/

    /**
     * 根据id查找
     *
     * @param indexName
     * @param typeName
     * @param id
     * @param tClass
     * @return
     * @throws Exception
     */
    public <T> T getById(String indexName, String typeName, String id, Class<T> tClass) {
        T obj = null;
        try {
            Get get = new Get.Builder(indexName, id).type(typeName).build();
            JestResult result = jestClient.execute(get);
            if (!result.isSucceeded()) {
                return null;
            }
            obj = result.getSourceAsObject(tClass);
        } catch (IOException e) {
            log.error("JestSearchUtil-getById异常:", e);
        }
        return obj;
    }

    /**
     * 根据ids查找
     *
     * @param indexName
     * @param typeName
     * @param ids
     * @param tClass
     * @return
     * @throws Exception
     */
    public <T> List<T> getListByIds(String indexName, String typeName, List<String> ids, Class<T> tClass) {
        List<T> obj = null;
        try {
            MultiGet build = new MultiGet.Builder.ById(indexName, typeName).addId(ids).build();
            JestResult result = jestClient.execute(build);
            if (!result.isSucceeded()) {
                return null;
            }
//            obj = result.getJsonString();
            obj = result.getSourceAsObjectList(tClass);
        } catch (IOException e) {
            log.error("JestSearchUtil-getListByIds异常:", e);
        }
        return obj;
    }

    /**
     * 根据查询语句，获取对应文档数量
     *
     * @param indexName
     * @param typeName
     * @return
     * @throws Exception
     */
    public Double count(String indexName, String typeName, String query) throws Exception {
        Count count = new Count.Builder()
                .addIndex(indexName)
                .addType(typeName)
                .query(query)
                .build();
        CountResult results = jestClient.execute(count);
        Double counts = results.getCount();
        return counts;
    }

    /**
     * 根据任意语句，执行查询
     *
     * @param indexName
     * @param typeName
     * @param query     任意语句
     * @throws Exception
     */
    public SearchResult query(String indexName, String typeName, String query) {
        SearchResult searchResult = null;
        try {
            Search search = new Search.Builder(query)
                    .addIndex(indexName)
                    .addType(typeName)
                    .build();
            searchResult = jestClient.execute(search);
        } catch (IOException e) {
            log.error("JestSearchUtil-query异常:", e);
        }
        return searchResult;
    }

    /**
     * 建议查询：通过completion查询匹配前缀查询
     * 仅返回文档的list集合
     */
    public List<String> suggest(String indexName, String typeName, String dimWords) {
        SearchResult searchResult = null;
        String suggestionName = "my-suggestion";
        Search suggest = new Search.Builder("{\n" +
                "  \"suggest\":{\n" +
                "    \"" + suggestionName + "\":{\n" +
                "      \"prefix\":\""+dimWords+"\",\n" +
                "      \"completion\":{\n" +
                "        \"field\":\"suggest\"\n" +
                "      }\n" +
                "    }\n" +
                "  }\n" +
                "}")
                .addIndex(indexName)
                .addType(typeName)
                .build();
        try {
            searchResult = jestClient.execute(suggest);
        } catch (IOException e) {
            log.error("JestSearchUtil-suggest异常:", e);
        }

        //解析SearchResult,只显示文档content内容
        List<String> contents = new ArrayList<>();
        String jsonString = searchResult.getJsonString();
        Map map1 = (Map) JSONObject.parse(jsonString);
        Map map2 = (Map) map1.get("suggest");
        List<Map> list1 = (List) map2.get("my-suggestion");
        for (Map map : list1) {
            List<Map> options = (List) map.get("options");
            for (Map option : options) {
                Map source = (Map) option.get("_source");
                String content = (String) source.get("content");
                contents.add(content);
            }
        }
        return contents;
    }
}
